import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib
matplotlib.use('agg')
import matplotlib.pyplot as plt
import pandas as pd
import json
from tmm import coh_tmm

# 使用tmm库添加材料
def load_material_data(file_name):
    data = pd.read_csv(file_name)
    wavelength = data["wl"].values  
    n = data["n"].values
    k = data["k"].values
    return wavelength, n, k

class SppSim:
    def __init__(self, fmin, fmax, nfreq, dW=0.5, Y=4, decay=1e-3, visualization=False):
        self.resolution = 100
        self.decay = decay
        self.visualization = visualization
        self.Y = Y
        self.dW = dW
        self.fmax = fmax
        self.fmin = fmin
        self.nfreq = nfreq
        self.fcen = (fmax + fmin) / 2
        self.df = (fmax - fmin)

    def simulate(self, h, dmp, thicknesses):
        wavelengths = np.linspace(0.28, 4, 600) 
        absorptions = []

        for wl in wavelengths:
            n_list = [1.0] 
            d_list = [np.inf]  

            for i, thickness in enumerate(thicknesses):
                if thickness == 0:
                    continue  
                if i % 2 == 0:
                    n_real = np.interp(wl, W_data[0], W_data[1])
                    k = np.interp(wl, W_data[0], W_data[2])
                    n = n_real + 1j * k
                else:
                    n_real = np.interp(wl, Al2O3_data[0], Al2O3_data[1])
                    k = np.interp(wl, Al2O3_data[0], Al2O3_data[2])
                    n = n_real + 1j * k
                n_list.append(n)
                d_list.append(thickness)  

            n_list.append(1.0) 
            d_list.append(np.inf)  

            # 设置偏振和入射角
            pol = 'p'
            inc_angle = 0

            try:
                # 调用coh_tmm计算反射和透射
                coh_tmm_data = coh_tmm(pol, n_list, d_list, inc_angle, wl)
                # 提取吸收率
                absorption = 1 - coh_tmm_data['R'] - coh_tmm_data['T']
                absorptions.append(absorption)
            except ValueError as e:
                print(f"coh_tmm error at wavelength {wl}: {e}")
                absorptions.append(0) 

        return np.array(absorptions)

class Spp:
    def __init__(self, fmin, fmax, nfreq, layers):
        self.sim = SppSim(fmin, fmax, nfreq)
        self.layers = layers
        self.pbounds = {
            'rh': (0.05, 0.3),   # 厚度限制为0.05到0.3
            'rdmp': (0, 2),      # 阻尼限制
        }
        # 为每一层材料设置厚度限制
        for i, (thickness, material) in enumerate(layers):
            if material == 'W':
                self.pbounds[f'x{i+1}'] = (0.05, 0.3)  # 对W层的厚度限制
            elif material == 'Al2O3':
                self.pbounds[f'x{i+1}'] = (0.1, 0.5)  # 对Al2O3层的厚度限制
            else:
                self.pbounds[f'x{i+1}'] = (0.01, 0.5)  # 其他材料的默认限制

    def evaluate(self, **params):
        h = params['rh']
        dmp = params['rdmp']
        thicknesses = [params[f'x{i+1}'] for i in range(len(self.layers))]
        absorption = self.sim.simulate(h, dmp, thicknesses)
        result = np.max(absorption)  
        print(f"Result: {result}, Params: {params}")

        with open("Xy.json", "a") as f:
            json.dump((float(result), params), f)
            f.write('\n')
        return result


# DNGO模型定义
class DNGO(nn.Module):
    def __init__(self, input_dim, hidden_dim=50):
        super(DNGO, self).__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        self.fc3 = nn.Linear(hidden_dim, 1)
        
    def forward(self, x):
        x = torch.tanh(self.fc1(x))
        x = torch.tanh(self.fc2(x))
        output = self.fc3(x)
        return output

# DNGO优化器
class DNGOptimizer:
    def __init__(self, f, pbounds, input_dim, hidden_dim=50, learning_rate=1e-3):
        self.f = f
        self.pbounds = pbounds
        self.input_dim = input_dim
        self.model = DNGO(input_dim, hidden_dim)
        self.optimizer = optim.Adam(self.model.parameters(), lr=learning_rate)
        self.loss_fn = nn.MSELoss()
        self.samples = []
        self.targets = []

    def maximize(self, init_points=5, n_iter=25):
        # 初始化样本点
        for _ in range(init_points):
            params = {k: np.random.uniform(*v) for k, v in self.pbounds.items()}
            result = self.f(**params)
            self.samples.append(list(params.values()))
            self.targets.append(result)

        # 开始训练和优化过程
        for _ in range(n_iter):
            self.train_model()
            next_params = self.suggest_next_params()
            next_result = self.f(**next_params)
            self.samples.append(list(next_params.values()))
            self.targets.append(next_result)
            print(f"Iteration {_+1}, Params: {next_params}, Result: {next_result}")

    def train_model(self, epochs=1000):
        x_train = torch.tensor(self.samples, dtype=torch.float32)
        y_train = torch.tensor(self.targets, dtype=torch.float32).view(-1, 1)
        self.model.train()

        for epoch in range(epochs):
            self.optimizer.zero_grad()
            predictions = self.model(x_train)
            loss = self.loss_fn(predictions, y_train)
            loss.backward()
            self.optimizer.step()

    def suggest_next_params(self):
        # 随机搜索最优参数
        best_score = float('-inf')
        best_params = None
        for _ in range(100):
            params = {k: np.random.uniform(*v) for k, v in self.pbounds.items()}
            params_list = torch.tensor(list(params.values()), dtype=torch.float32)
            with torch.no_grad():
                pred = self.model(params_list).item()
            if pred > best_score:
                best_score = pred
                best_params = params
        return best_params

if __name__ == '__main__':
    # 读取材料数据
    W_data = load_material_data('/data2/belly/Project/交付/run/W.csv')
    Al2O3_data = load_material_data('/data2/belly/Project/交付/run/Al2O3.csv')

    # 定义材料层结构
    layers = [
        (0.1, 'W'), (0.1, 'Al2O3'),
        (0.1, 'W'), (0.1, 'Al2O3'),
        (0.1, 'W'), (0.1, 'Al2O3'),
        (0.1, 'W'), (0.1, 'Al2O3'),
        (0.1, 'W')  # 最后一层W
    ]  # 共9层结构

    sim = Spp(0.28, 4, 600, layers)

    # 定义目标函数
    def objective(**params):
        return sim.evaluate(**params)

    # 初始化 DNGO 优化器
    optimizer = DNGOptimizer(
        f=objective,
        pbounds=sim.pbounds,
        input_dim=len(sim.pbounds)
    )

    optimizer.maximize(init_points=5, n_iter=50)

    # 找到最大效率的索引
    max_index = np.argmax(optimizer.targets)
    max_efficiency = optimizer.targets[max_index]
    max_params = optimizer.samples[max_index]

    # 保存最大效率到文件
    with open("max_efficiency.json", "w") as f:
        json.dump({
            "max_efficiency": float(max_efficiency),
            "thicknesses": max_params
        }, f, indent=4)

    print(f"最大效率: {max_efficiency}")

    # 生成波长-吸光率图片
    thicknesses_best = max_params
    absorption_best = sim.sim.simulate(max_params[0], max_params[1], thicknesses_best)
    wavelengths = np.linspace(0.28, 4, 600)

    plt.figure(figsize=(10, 6))
    plt.plot(wavelengths, absorption_best, label='最大效率吸光率')
    plt.xlabel('波长')
    plt.ylabel('吸光率')
    plt.title('波长-吸光率曲线（最大效率）')
    plt.legend()
    plt.grid(True)
    plt.savefig('wavelength_absorption.png')
    plt.close()

    # 动态跟踪最大效率
    max_eff_list = []
    current_max = -np.inf
    for i, res in enumerate(optimizer.targets):
        current_max = max(current_max, res)  # 跟踪迄今为止的最大效率
        max_eff_list.append(current_max)

    # 绘制迭代次数-最大效率图
    plt.figure(figsize=(10, 6))
    plt.plot(range(1, len(max_eff_list) + 1), max_eff_list, marker='o')
    plt.xlabel('计算次数')
    plt.ylabel('最大效率')
    plt.title('计算次数-最大效率关系')
    plt.grid(True)
    plt.savefig('iterations_max_efficiency.png')
    plt.close()

    print("Optimization complete. Results saved as follows:")
    print("1. 最大效率: max_efficiency.json")
    print("2. 波长-吸光率图片: wavelength_absorption.png")
    print("3. 计算次数-最大效率图片: iterations_max_efficiency.png")
